import wfdb
import numpy as np
import csv
from feature_extraction import main_feature

from myutils.NcUtils import fliter
from myutils.beatsLenUtils import extract_variable_length_beats, normalize_beat_length
from myutils.dataTypeUtils import ECG_R_list, ECGSignal, AAMI_MIT, AAMI

# warnings.simplefilter('always')

# TODO: 病人间模式，变长度，自适应插值,医学特征

# 所有数据
# file_name = ['100', '101', '102', '103', '104', '105', '106', '107',
#              '108', '109', '111', '112', '113', '114', '115', '116',
#              '117', '118', '119', '121', '122', '123', '124', '200',
#              '201', '202', '203', '205', '207', '208', '209', '210',
#              '212', '213', '214', '215', '217', '219', '220', '221',
#              '222', '223', '228', '230', '231', '232', '233', '234']


# 北交论文中的病人间分类
# 训练集
# file_name = ['101', '106', '108', '109', '112', '114', '115', '116',
#              '118', '119', '122', '124', '201', '203', '205', '207',
#              '208', '209', '215', '220', '223', '230']
# 测试集
file_name = ['100', '103', '105', '111', '113', '117', '121', '123',
             '200', '202', '210', '212', '213', '214', '219', '221',
             '222', '228', '231', '232', '233', '234']

# MILL+V1导联，共40个人，按如下规则划分，
# 带导联的病人间
# file_name = ['101', '106', '108', '109', '112', '115', '116', '118', '119', '122',
#              '201', '203', '205', '207', '208', '209', '215', '220', '223', '230']
# file_name = ['105', '107', '111', '113', '121', '200', '202', '210', '212', '213',
#              '214', '217', '219', '221', '222', '228', '231', '232',  '234']
# '233',
# file_name = ['101', '112', '113', '115', '122', '212', '220', '222', '232', '111', '121', '209', '230', '231', '234',
#              '118', '105', '107', '116', '217', '207', '228', '221', '119', '106', '213', '208']
# file_name = ['202', '219', '215', '214', '203', '108', '109', '201', '200', '210', '205', '233', '223']

for F_name in file_name:
    print(f'正在处理数据{F_name}')
    # wfdb.rdann 用于从注释文件中读取注释信息。这些注释通常包含心拍标记、心律失常类型等重要信息。
    # record_name：记录的名称，通常是文件名的前缀。
    # extension：注释文件的扩展名，常见的有'atr' 和'qrs'等。
    # sampfrom：从哪个样本开始读取注释，默认为0。
    # sampto：到哪个样本结束读取注释，默认为None，表示读取所有注释。
    # shift_samps：是否将样本索引加上偏移量，默认为False。

    # wfdb.rdann 返回一个 wfdb.Annotation 对象，该对象包含了以下属性：
    # sample：注释的时间点（以样本索引表示）。
    # symbol：注释符号，表示心拍类型或事件。
    # aux_note：辅助注释，可能为空。
    # fs：采样频率。
    # counter：计数器，用于跟踪注释的数量。
    signal_annotation = wfdb.rdann(f'../data/mit-bih-arrhythmia-database-1.0.0/{F_name}', "atr", sampfrom=0,
                                   sampto=650000)
    # 参数说明
    #  record_name：记录的名称，通常是文件名的前缀。
    #  channels：指定要读取的通道列表，默认为所有通道。例如，[0, 1]
    # 表示读取第一个和第二个通道。
    # sampfrom：从哪个样本开始读取，默认为0。
    # sampto：到哪个样本结束读取，默认为None，表示读取所有样本。
    # physical：是否将信号转换为物理单位。如果为True，则信号将以物理单位返回；如果为False，则信号将以数字单位返回。
    # 返回值
    # wfdb.rdrecord 返回一个 wfdb.Record 对象，该对象包含了以下属性：
    # p_signal：以物理单位表示的信号数据（如果physical = True）。
    # d_signal：以数字单位表示的信号数据（如果physical = False）。
    # fs：采样频率。
    # units：每个通道的单位。
    # sig_name：每个通道的名称。
    # base_date：记录的基准日期。
    # base_time：记录的基准时间。
    #  comments：记录的注释信息。
    record = wfdb.rdrecord(f'../data/mit-bih-arrhythmia-database-1.0.0/{F_name}', sampfrom=0, sampto=650000,
                           physical=True,
                           channels=[0])
    # if record.sig_name != ['MLII', 'V1']:
    #     print("{F_name}类型错误")
    #     continue
    # ecg = record.p_signal.flatten()
    arr_transposed = record.p_signal.T  # 现在形状是 (2, 32)
    # arr1 = arr_transposed[0]
    # arr2 = arr_transposed[1]

    for i in range(1):
        arr1 = arr_transposed[i]
        r_peaks = signal_annotation.sample
        arr1 = fliter(arr1)

        areaVector, featureVector, amplitudeVector, RR_interval = main_feature(arr1, r_peaks, 360)
        variable_length_beats, intervals, beat_symbol = extract_variable_length_beats(arr1, r_peaks,
                                                                                      signal_annotation.symbol)
        normalized_beats = normalize_beat_length(variable_length_beats, target_length=288)
        if len(normalized_beats) != len(variable_length_beats):
            print("数据量不一致")
            exit(0)

        # 获取心拍种类
        Label_kind = list(set(beat_symbol))
        # 读取每一种R点在信号中的位置
        for k in Label_kind:
            if k in ECG_R_list:
                index = [i for i, x in enumerate(beat_symbol) if x == k]
                for site in index:
                    ECG_signal_feature = np.concatenate((amplitudeVector[site], featureVector[site]))
                    ECG_signal_feature = np.concatenate((ECG_signal_feature, areaVector[site]))
                    ECG_signal_feature = np.concatenate((ECG_signal_feature, normalized_beats[site]))
                    ECGSignal[str(k)].append(ECG_signal_feature)

# 打印种类
for key, value in ECGSignal.items():
    print(f'{key} = {len(value)}')

for ECG_key, ECG_value in ECGSignal.items():
    for AAMI_MIT_key, AAMI_MIT_value in AAMI_MIT.items():
        if ECG_key in AAMI_MIT_value:
            AAMI[AAMI_MIT_key].extend(ECG_value)

# 5分类并保存到CSV格式的文件里
for key, value in AAMI.items():
    with open(f'../data/mit/brn52_268/{key}-test-251020.csv', 'w', newline='\n') as f:
        writer = csv.writer(f)
        # 将列表的每条数据依次写入csv文件， 并以逗号分隔
        # 传入的数据为列表中嵌套列表或元组，每一个列表或元组为每一行的数据
        writer.writerows(value)
